package com.ruoyi.aidesign.websocket;

import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesis;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisParam;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisResult;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.exception.UploadFileException;
import com.alibaba.dashscope.utils.JsonUtils;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.ruoyi.aidesign.model.DialogueMessage;
import com.ruoyi.aidesign.model.HouseDesignRequest;
import com.ruoyi.aidesign.model.MinioImageProcessor;
import jakarta.websocket.*;
import jakarta.websocket.server.ServerEndpoint;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.*;

import com.alibaba.dashscope.aigc.multimodalconversation.MultiModalConversation;
import com.alibaba.dashscope.aigc.multimodalconversation.MultiModalConversationParam;
import com.alibaba.dashscope.aigc.multimodalconversation.MultiModalConversationResult;
import com.alibaba.dashscope.common.MultiModalMessage;
import com.alibaba.dashscope.common.Role;


import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 设计对话WebSocket处理器
 */
@Slf4j
@Component
@ServerEndpoint("/aidesign/dialogue")
public class DesignDialogueWebSocket {

    /**
     * 当前在线连接数
     */
    private static AtomicInteger onlineCount = new AtomicInteger(0);

    /**
     * 用于存放所有在线客户端
     */
    private static Map<String, DesignDialogueWebSocket> clients = new ConcurrentHashMap<>();

    /**
     * 客户端对应的会话
     */
    private Session session;

    /**
     * 客户端ID
     */
    private String clientId;

    /**
     * JSON工具类
     */
    private static ObjectMapper objectMapper = new ObjectMapper();

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session) {
        this.session = session;
        this.clientId = session.getId();
        clients.put(clientId, this);

        // 添加在线人数
        int count = onlineCount.incrementAndGet();
        log.info("有新连接加入，当前在线人数为：{}", count);

        // 发送连接成功消息
        sendMessage(DialogueMessage.builder()
                .type(DialogueMessage.MessageType.CONNECTED)
                .content("")
                .timestamp(System.currentTimeMillis())
                .source("system")
                .status("success")
                .build());
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public void onClose() {
        clients.remove(clientId);
        int count = onlineCount.decrementAndGet();
        log.info("有一连接关闭，当前在线人数为：{}", count);
    }

    /**
     * 收到客户端消息后调用的方法
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        log.info("收到来自客户端{}的消息：{}", clientId, message);

        try {
            // 解析请求数据
            HouseDesignRequest request = objectMapper.readValue(message, HouseDesignRequest.class);

            // 发送确认消息
            sendMessage(DialogueMessage.builder()
                    .type(DialogueMessage.MessageType.RESPONSE)
                    .content("已收到您的设计请求，正在为您生成方案...")
                    .timestamp(System.currentTimeMillis())
                    .source("system")
                    .status("processing")
                    .build());

            // 模拟处理过程
            new Thread(() -> {
                try {

                    // 组装户型和风格信息
                    StringBuilder resultMessage = new StringBuilder();

                    Map<String, String> roomPic = null;
                    try {
                        roomPic = simpleMultiModalConversationCall(request.getHouseValue().getHousePic(), request.getStyleValue().getImage());
                        for (String s : roomPic.keySet()) {
                            System.out.println(s);
                        }
                    } catch (ApiException | NoApiKeyException | UploadFileException e) {
                        log.error("调用图片生成API发生错误", e);
                        sendErrorMessage("生成图片时发生错误: " + e.getMessage());
                        return;
                    }

                    // 发送最终结果，包含房间图片数据
                    DialogueMessage resultDialogueMessage = DialogueMessage.builder()
                            .type(DialogueMessage.MessageType.RESPONSE)
                            .content(resultMessage.toString())
                            .timestamp(System.currentTimeMillis())
                            .source("ai")
                            .status("completed")
                            .build();

                    // 将roomPic添加到消息中
                    resultDialogueMessage.setRoomPics(roomPic);
                    sendMessage(resultDialogueMessage);

                } catch (Exception e) {
                    log.error("处理过程被中断", e);
                    sendErrorMessage("生成过程被中断，请重试");
                }
            }).start();

        } catch (Exception e) {
            log.error("处理消息时发生错误", e);
            sendErrorMessage("无法处理您的请求，请检查输入格式");
        }
    }

    /**
     * 发送进度消息
     */
    private void sendProgressMessage(String message, int progress) {
        sendMessage(DialogueMessage.builder()
                .type(DialogueMessage.MessageType.PROGRESS)
                .content(message)
                .timestamp(System.currentTimeMillis())
                .source("system")
                .status(String.valueOf(progress))
                .build());
    }

    /**
     * 发送错误消息
     */
    private void sendErrorMessage(String errorMessage) {
        sendMessage(DialogueMessage.builder()
                .type(DialogueMessage.MessageType.ERROR)
                .content(errorMessage)
                .timestamp(System.currentTimeMillis())
                .source("system")
                .status("error")
                .build());
    }

    /**
     * 发生错误时调用
     */
    @OnError
    public void onError(Session session, Throwable error) {
        log.error("发生错误", error);
        sendErrorMessage("连接发生错误，请刷新页面重试");
    }

    /**
     * 发送消息
     */
    private void sendMessage(DialogueMessage message) {
        try {
            this.session.getBasicRemote().sendText(objectMapper.writeValueAsString(message));
        } catch (IOException e) {
            log.error("发送消息时发生错误", e);
        }
    }

    /**
     * 获取当前在线人数
     */
    public static int getOnlineCount() {
        return onlineCount.get();
    }


    // 图片转文字

    public static Map<String,String> simpleMultiModalConversationCall(String house, String style) throws ApiException, NoApiKeyException, UploadFileException {
        MultiModalConversation conv = new MultiModalConversation();
        // 如果使用本地图像，请导入 import java.util.HashMap;，再为函数添加【String localPath】参数，并释放下面注释
        // String filePath = "file://"+localPath;
        MultiModalMessage systemMessage = MultiModalMessage.builder().role(Role.SYSTEM.getValue())
                .content(Arrays.asList(
                        Collections.singletonMap("text", "你是一个AI全屋设计智慧机器人"))).build();
        MultiModalMessage userMessage = MultiModalMessage.builder().role(Role.USER.getValue())
                .content(Arrays.asList(
                        // 第一张图像链接
                        Collections.singletonMap("image", house),
                        // 如果使用本地图像，请并释放下面注释
                        // new HashMap<String, Object>(){{put("image", filePath);}},
                        // 第二张图像链接
                        Collections.singletonMap("image", style),
                        // 第三张图像链接
//                        Collections.singletonMap("image", "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/hbygyo/rabbit.jpg"),
                        Collections.singletonMap("text", "请根据以上的图片，生成每个房间最详细的装修方案,每个房间的方案输出完成后利用'----',在每个房间的风格中都写上北欧风格"))).build();
        MultiModalConversationParam param = MultiModalConversationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                .model("qwen-vl-max-latest")
                .messages(Arrays.asList(systemMessage, userMessage))
                .build();
        MultiModalConversationResult result = conv.call(param);
        String args = result.getOutput().getChoices().get(0).getMessage().getContent().get(0).get("text").toString();
        System.out.println(args);
        System.out.println(".....");
        // 获取房间装修方案
        List<String> roomPlans = splitRoomPlans(args);
        for (int i = 0; i < roomPlans.size(); i++) {
            System.out.println("房间 " + (i + 1) + " 装修方案：");
            System.out.println(roomPlans.get(i));
            System.out.println("-------------------");
        }
        if (roomPlans.size() > 0) {
            roomPlans.remove(roomPlans.size() - 1);
        }
        // 获取房间名

        List<String> roomNames = extractRoomNames(args);
        String[] roomNameArray = roomNames.toArray(new String[0]);
        // 生成的所有图片
        List<String> Pics = new ArrayList<>();
        for (String part : roomPlans) {
            System.out.println(part);
            try{
                System.out.println("this");
                Pics.add(basicCall(part));
                //listTask();
            }catch(ApiException|NoApiKeyException e){
                System.out.println(e.getMessage());
            }
        }
        Map<String, String> roomPics = new HashMap<>();
        for (int i = 0; i < roomNameArray.length; i++) {
            roomPics.put(roomNameArray[i], Pics.get(i));
        }
        return roomPics;
    }

    // 文字转图片
    public static String basicCall(String part) throws ApiException, NoApiKeyException {
        ImageSynthesisParam param =
                ImageSynthesisParam.builder()
                        .apiKey(System.getenv("DASHSCOPE_API_KEY"))
                        .model("wanx2.1-t2i-turbo")
                        .prompt(part)
                        .n(1)
                        .size("720*720")
                        .build();
        ImageSynthesis imageSynthesis = new ImageSynthesis();
        ImageSynthesisResult result = null;
        try {
            System.out.println("---sync call, please wait a moment----");
            result = imageSynthesis.call(param);
        } catch (ApiException | NoApiKeyException e){
            throw new RuntimeException(e.getMessage());
        }

        // 将结果转为JSON字符串
        String resultJson = JsonUtils.toJson(result);
        System.out.println(resultJson);

        // 处理图片URL并上传到MinIO
        String minioUrl = MinioImageProcessor.processApiResult(resultJson);
        if (minioUrl != null) {
            return  minioUrl;
        } else {
            return "";
        }
    }
    public static List<String> splitRoomPlans(String allPlans) {
        List<String> plans = new ArrayList<>();
        String[] parts = allPlans.split("---");
        for (String part : parts) {
            part = part.trim();
            if (!part.isEmpty()) {
                plans.add(part);
            }
        }
        return plans;
    }

    public static List<String> extractRoomNames(String text) {
        List<String> roomNames = new ArrayList<>();
        Pattern pattern = Pattern.compile("### \\d+\\. \\*\\*(.*?)装修方案\\*\\*");
        Matcher matcher = pattern.matcher(text);

        while (matcher.find()) {
            roomNames.add(matcher.group(1));
        }
        return roomNames;
    }
}
